from typing import Callable, Optional
import logging
import asyncio
from fastapi import WebSocket
from starlette import websockets
from gurren.app import app
from gurren.providers import gen_sin_distortion

log = logging.getLogger(__name__)


def websocket_close(fn: Optional[Callable] = None):
    inner_fn = fn    
    async def closer(ws: WebSocket):
        assert inner_fn is not None
        await inner_fn(ws)
        if ws.client_state != websockets.WebSocketState.DISCONNECTED:
            try:
                await ws.close()
            except Exception as ex:
                log.debug('websocket_close err=%s', ex)    

    def decorator(fn: Callable):
        nonlocal inner_fn
        inner_fn = fn
        return closer

    return closer if fn else decorator

@app.websocket("/ws/sin")
@websocket_close()
async def sin_endpoint(websocket: WebSocket):

    async def send_data():
        async for r in gen_sin_distortion():
            log.debug('WS r=%s cookie=%s', r, websocket.cookies.get("x-gurren-captcha"))
            await websocket.send_json(r)
    
    async def read_data():
        while True:
            ack = await websocket.receive_text()
            if ack.upper() == 'STOP':
                log.debug('WS client asks to stop')
                break
            await websocket.send_text('You said: ' + ack)
        return 

    def stop_tasks(*_args):
        log.debug('stop tasks=%s', ", ".join(f'<{t.get_name()} {t._state}>' for t in tasks))
        for t in tasks:
            if not t.done():
                log.debug('WS cancel task %s', t)
                t.cancel()
            t.remove_done_callback(stop_tasks)

    await websocket.accept()    
    if not websocket.cookies.get("x-gurren-captcha"):
        await websocket.send_json({'no': 'cookie'})
        return
    await websocket.send_json({'hello': 1})

    tasks = (asyncio.create_task(send_data(), name='send_data'),
             asyncio.create_task(read_data(), name='read_data'))
    for t in tasks:
        t.add_done_callback(stop_tasks)

    tasks_res = await asyncio.gather(*tasks, return_exceptions=True)
    log.info('WS client end %s', tasks_res)
